import org.junit.Assert;
import org.testng.annotations.Test;

import java.awt.*;
import java.awt.image.BufferedImage;
import java.io.IOException;
import java.net.HttpURLConnection;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;

public class LoadTestIT {
    public static final int LATCH_SIZE = 5;
    public static final int PHOTOS_PER_THREAD = 3;
    static CountDownLatch latch = new CountDownLatch(LATCH_SIZE);// jbossews-vitaliykuzmich.rhcloud.com localhost:8080/ROOT
    static final String UPLOAD_URL = "http://" + "jbossews-vitaliykuzmich.rhcloud.com" + "/upload";
    static final String charset = "UTF-8";
    static ConcurrentHashMap<Long, Long> times = new ConcurrentHashMap<>();
    static AtomicBoolean isFirst = new AtomicBoolean(false);

    @Test(invocationCount = LATCH_SIZE, threadPoolSize = LATCH_SIZE)
    public void loadTest() throws IOException, InterruptedException {
        System.out.println("Using server url:" + UPLOAD_URL);
        ArrayList<BufferedImage> arr = new ArrayList<>(PHOTOS_PER_THREAD);
        for (int i = 0; i < PHOTOS_PER_THREAD; i++) {
            arr.add(fillImageRandomly(new BufferedImage(300, 200, BufferedImage.TYPE_INT_RGB)));
        }

        Long begin = System.currentTimeMillis();
        MultipartUtility multipart = new MultipartUtility(UPLOAD_URL, charset);

        multipart.addFormField("w", "150");
        multipart.addFormField("h", "200");
        latch.countDown();

        latch.await();

        for (BufferedImage image : arr) {
            multipart.addFilePart("file", fillImageRandomly(image));
        }

        List<String> response = multipart.finish();
        times.put(Thread.currentThread().getId(), (Long) System.currentTimeMillis() - begin);


        System.out.println("SERVER REPLIED:");
        for (String line : response) {
            System.out.println(line);
        }
        if (!isFirst.getAndSet(true)) {
            Long total = 0l;
            for (Map.Entry<Long, Long> entry : times.entrySet()) {
                total += entry.getValue();
            }
            System.out.println();
            System.out.println();
            System.out.println("Threads count " + LATCH_SIZE);
            System.out.println("Photos per thread " + PHOTOS_PER_THREAD);
            System.out.println("Avg upload time " + total / LATCH_SIZE+" ms");
        }
        Assert.assertEquals(HttpURLConnection.HTTP_OK, multipart.getReturnedStatus());
    }


    private BufferedImage fillImageRandomly(BufferedImage img) {
        Color col = null;
        for (int i = 0; i < img.getWidth(); i++)
            for (int j = 0; j < img.getHeight(); j++) {
                col = new Color((int) (Math.random() * 255), (int) (Math.random() * 255), (int) (Math.random() * 255));
                img.setRGB(i, j, col.getRGB());
            }
        return img;
    }

}
